# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The Lance Authors

import os
import threading

import numpy as np
import pyarrow as pa
import pyarrow.compute as pc
import pyarrow.parquet as pq
import pytest
from lance.file import LanceFileReader, LanceFileSession, LanceFileWriter


def test_file_read_projection(tmp_path):
    table = pa.table({"a": [1, 2, 3], "b": [4, 5, 6]})
    path = tmp_path / "foo.lance"
    with LanceFileWriter(str(path)) as writer:
        writer.write_batch(table)

    reader = LanceFileReader(str(path), columns=["a"])
    assert reader.read_all().to_table() == table.select("a")


def test_file_writer(tmp_path):
    path = tmp_path / "foo.lance"
    schema = pa.schema([pa.field("a", pa.int64())])
    with LanceFileWriter(str(path), schema) as writer:
        writer.write_batch(pa.table({"a": [1, 2, 3]}))
    reader = LanceFileReader(str(path))
    metadata = reader.metadata()
    assert metadata.num_rows == 3


def test_write_no_schema(tmp_path):
    path = tmp_path / "foo.lance"
    with LanceFileWriter(str(path)) as writer:
        writer.write_batch(pa.table({"a": [1, 2, 3]}))
    reader = LanceFileReader(str(path))
    assert reader.read_all().to_table() == pa.table({"a": [1, 2, 3]})


def test_no_schema_no_data(tmp_path):
    path = tmp_path / "foo.lance"
    with pytest.raises(
        ValueError, match="Schema is unknown and file cannot be created"
    ):
        with LanceFileWriter(str(path)) as _:
            pass


def test_schema_only(tmp_path):
    path = tmp_path / "foo.lance"
    schema = pa.schema([pa.field("a", pa.int64())])
    with LanceFileWriter(str(path), schema=schema) as _:
        pass
    reader = LanceFileReader(str(path))
    assert reader.metadata().schema == schema


def test_write_with_max_page_bytes(tmp_path):
    path = tmp_path / "foo.lance"
    schema = pa.schema([pa.field("a", pa.int64())])
    for version in ["2.0", "2.1"]:
        with LanceFileWriter(
            str(path), schema, max_page_bytes=1, version=version
        ) as writer:
            writer.write_batch(pa.table({"a": [1, 2, 3]}))
        reader = LanceFileReader(str(path))
        # Only 2.0 splits large pages on write.   In 2.1+ we split on read.
        expected_pages = 3 if version == "2.0" else 1
        assert len(reader.metadata().columns[0].pages) == expected_pages


def test_aborted_write(tmp_path):
    path = tmp_path / "foo.lance"
    schema = pa.schema([pa.field("a", pa.int64())])
    writer = LanceFileWriter(str(path), schema)
    writer.write_batch(pa.table({"a": [1, 2, 3]}))
    del writer
    assert not path.exists()


def test_multiple_close(tmp_path):
    path = tmp_path / "foo.lance"
    schema = pa.schema([pa.field("a", pa.int64())])
    writer = LanceFileWriter(str(path), schema)
    writer.write_batch(pa.table({"a": [1, 2, 3]}))
    writer.close()
    writer.close()


def test_version(tmp_path):
    path = tmp_path / "foo.lance"
    schema = pa.schema([pa.field("a", pa.int64())])

    with LanceFileWriter(str(path), schema, version="2.0") as writer:
        writer.write_batch(pa.table({"a": [1, 2, 3]}))
    reader = LanceFileReader(str(path))
    metadata = reader.metadata()
    assert metadata.major_version == 0
    assert metadata.minor_version == 3

    path = tmp_path / "foo2.lance"
    with LanceFileWriter(str(path), schema, version="2.1") as writer:
        writer.write_batch(pa.table({"a": [1, 2, 3]}))
    reader = LanceFileReader(str(path))
    metadata = reader.metadata()
    assert metadata.major_version == 2
    assert metadata.minor_version == 1


def test_take(tmp_path):
    path = tmp_path / "foo.lance"
    schema = pa.schema([pa.field("a", pa.int64())])
    writer = LanceFileWriter(str(path), schema)
    writer.write_batch(pa.table({"a": [i for i in range(100)]}))
    writer.close()

    reader = LanceFileReader(str(path))
    # Can't read out of range
    with pytest.raises(ValueError):
        reader.take_rows([0, 100]).to_table()

    table = reader.take_rows([0, 77, 83]).to_table()
    assert table == pa.table({"a": [0, 77, 83]})


def test_num_rows(tmp_path):
    path = tmp_path / "foo.lance"
    schema = pa.schema([pa.field("a", pa.int64())])
    writer = LanceFileWriter(str(path), schema)
    writer.write_batch(pa.table({"a": [i for i in range(100)]}))
    writer.close()

    reader = LanceFileReader(str(path))
    assert reader.num_rows() == 100


def check_round_trip(tmp_path, table):
    path = tmp_path / "foo.lance"
    with LanceFileWriter(str(path), table.schema) as writer:
        for batch in table.to_batches():
            writer.write_batch(pa.table(batch))
    reader = LanceFileReader(str(path))
    result = reader.read_all().to_table()
    assert result == table


def test_different_types(tmp_path):
    dict_values = pa.array(["foo", "bar", "baz"], pa.string())
    dict_indices = pa.array([2, 1, 0], pa.uint8())

    check_round_trip(
        tmp_path,
        pa.table(
            {
                "large_string": pa.array(["foo", "bar", "baz"], pa.large_string()),
                "large_binary": pa.array([b"foo", b"bar", b"baz"], pa.large_binary()),
                "dict_string": pa.DictionaryArray.from_arrays(
                    dict_indices, dict_values
                ),
            }
        ),
    )


def test_with_nulls(tmp_path):
    check_round_trip(
        tmp_path,
        pa.table(
            {
                "some_null_1": pa.array([1, 2, None], pa.int64()),
                "some_null_2": pa.array([None, None, 3], pa.int64()),
                "nullable_list": pa.array(
                    [[1, 2], None, [None, 3]], pa.list_(pa.int64())
                ),
                "nullable_fsl": pa.array(
                    [[1, 2], None, [None, 3]], pa.list_(pa.int64(), 2)
                ),
                "all_null": pa.array([None, None, None], pa.int64()),
                "null_strings": pa.array([None, "foo", None], pa.string()),
            }
        ),
    )


def test_batch_sizes(tmp_path):
    # Need a big string so there aren't too many rows per page because we
    # want to test different page sizes:
    #  - batch that spans multiple pages (including more than 2)
    #  - batch that is smaller than a page (including much smaller)
    my_str = b"0" * 299593

    data = [[my_str] for _ in range(1009)]
    tab = pa.table({"val": data})

    path = str(tmp_path / "foo.lance")
    with LanceFileWriter(path) as writer:
        writer.write_batch(tab)

    reader = LanceFileReader(path)

    for batch_size in range(10, 1050, 10):
        reader.read_all(batch_size=batch_size).to_table()


def test_round_trip(tmp_path):
    path = tmp_path / "foo.lance"
    schema = pa.schema([pa.field("a", pa.int64())])
    data = pa.table({"a": [1, 2, 3]})
    with LanceFileWriter(str(path), schema) as writer:
        writer.write_batch(data)
    reader = LanceFileReader(str(path))
    result = reader.read_all().to_table()
    assert result == data

    # TODO: Currently fails, need to fix reader
    # result = reader.read_range(1, 1).to_table()
    # assert result == pa.table({"a": [2]})

    # TODO: Test reading invalid ranges
    # TODO: Test invalid batch sizes


def test_metadata(tmp_path):
    path = tmp_path / "foo.lance"
    schema = pa.schema([pa.field("a", pa.int64())])
    data = pa.table({"a": [1, 2, 3]})
    with LanceFileWriter(str(path), schema) as writer:
        writer.write_batch(data)
    reader = LanceFileReader(str(path))
    metadata = reader.metadata()

    assert metadata.schema == schema
    assert metadata.num_rows == 3
    assert metadata.num_global_buffer_bytes > 0
    assert metadata.num_column_metadata_bytes > 0
    assert metadata.num_data_bytes > 0
    assert len(metadata.columns) == 1

    column = metadata.columns[0]
    assert len(column.column_buffers) == 0
    assert len(column.pages) == 1

    page = column.pages[0]
    assert len(page.buffers) == 1
    assert page.buffers[0].position == 0
    assert page.buffers[0].size == 24

    assert len(page.encoding) > 0


def test_file_stat(tmp_path):
    path = tmp_path / "foo.lance"
    schema = pa.schema(
        [pa.field("a", pa.int64()), pa.field("b", pa.list_(pa.float64(), 8))]
    )

    num_rows = 1_000_000

    data1 = pa.array(range(num_rows))

    # Create a fixed-size list of float64 with dimension 8
    fixed_size_list = [np.random.rand(8).tolist() for _ in range(num_rows)]
    data2 = pa.array(fixed_size_list, type=pa.list_(pa.float64(), 8))

    with LanceFileWriter(str(path), schema) as writer:
        writer.write_batch(pa.table({"a": data1, "b": data2}))
    reader = LanceFileReader(str(path))
    file_stat = reader.file_statistics()

    assert len(file_stat.columns) == 2

    assert file_stat.columns[0].num_pages == 1
    assert file_stat.columns[0].size_bytes <= 8_000_000

    # 2 pages on 2.0, 1 page in 2.1+
    assert file_stat.columns[1].num_pages <= 2
    # Slightly larger than 64MiB because of padding, chunk overhead, etc.
    assert file_stat.columns[1].size_bytes <= 64_200_000


def test_round_trip_parquet(tmp_path):
    pq_path = tmp_path / "foo.parquet"
    table = pa.table({"int": [1, 2], "list_str": [["x", "yz", "abc"], ["foo", "bar"]]})
    pq.write_table(table, str(pq_path))
    table = pq.read_table(str(pq_path))

    lance_path = tmp_path / "foo.lance"
    with LanceFileWriter(str(lance_path)) as writer:
        writer.write_batch(table)

    reader = LanceFileReader(str(lance_path))
    round_tripped = reader.read_all().to_table()
    assert round_tripped == table


def test_write_read_with_session(tmp_path):
    session = LanceFileSession(tmp_path)
    with session.open_writer("foo.lance") as writer:
        writer.write_batch(pa.table({"a": [1, 2, 3]}))

    with session.open_writer("bar.lance") as writer:
        writer.write_batch(pa.table({"a": [4, 5, 6]}))

    reader = session.open_reader("foo.lance")
    assert reader.read_all().to_table() == pa.table({"a": [1, 2, 3]})

    reader = session.open_reader("bar.lance")
    assert reader.read_all().to_table() == pa.table({"a": [4, 5, 6]})


def test_list_field_name(tmp_path):
    weird_field = pa.field("why does this name even exist", pa.string())
    weird_string_type = pa.list_(weird_field)
    schema = pa.schema([pa.field("list_str", weird_string_type)])
    table = pa.table({"list_str": [["x", "yz", "abc"], ["foo", "bar"]]}, schema=schema)

    path = tmp_path / "foo.lance"
    with LanceFileWriter(str(path)) as writer:
        writer.write_batch(table)

    reader = LanceFileReader(str(path))
    round_tripped = reader.read_all().to_table()

    assert round_tripped == table
    assert round_tripped.schema.field("list_str").type == weird_string_type


def test_field_meta(tmp_path):
    schema = pa.schema(
        [
            pa.field("primitive", pa.int64(), metadata={"foo": "bar"}),
            pa.field(
                "list",
                pa.list_(pa.field("item", pa.int64(), metadata={"list": "yes"})),
                metadata={"foo": "baz"},
            ),
            pa.field(
                "struct",
                pa.struct([pa.field("a", pa.int64(), metadata={"struct": "yes"})]),
                metadata={"foo": "qux"},
            ),
        ]
    )
    table = pa.table(
        {
            "primitive": [1, 2, 3],
            "list": [[1, 2], [3, 4], [5, 6]],
            "struct": [{"a": 1}, {"a": 2}, {"a": 3}],
        },
        schema=schema,
    )

    with LanceFileWriter(str(tmp_path / "foo.lance")) as writer:
        writer.write_batch(table)

    reader = LanceFileReader(str(tmp_path / "foo.lance"))
    round_tripped = reader.read_all().to_table()

    assert round_tripped == table


def test_dictionary(tmp_path):
    # Basic round trip
    dictionary = pa.array(["foo", "bar", "baz"], pa.string())
    indices = pa.array([0, 1, 2, 0, 1, 2], pa.int32())
    dict_arr = pa.DictionaryArray.from_arrays(indices, dictionary)

    def round_trip(arr):
        table = pa.table({"dict": arr})

        path = tmp_path / "foo.lance"
        with LanceFileWriter(str(path)) as writer:
            writer.write_batch(table)

        reader = LanceFileReader(str(path))
        table2 = reader.read_all().to_table()
        return table2.column("dict").chunk(0)

    round_tripped = round_trip(dict_arr)

    assert round_tripped == dict_arr
    assert round_tripped.type == dict_arr.type

    # Dictionary that doesn't use all values
    dictionary = pa.array(["foo", "bar", "baz"], pa.string())
    indices = pa.array([0, 0, 1, 1], pa.int32())
    dict_arr = pa.DictionaryArray.from_arrays(indices, dictionary)

    round_tripped = round_trip(dict_arr)

    assert round_tripped.dictionary == dictionary

    # different indices types
    dictionary = pa.array(["foo", "bar", "baz"], pa.string())
    for data_type in [
        pa.uint8(),
        pa.uint16(),
        pa.uint32(),
        pa.uint64(),
        pa.int8(),
        pa.int16(),
        pa.int32(),
        pa.int64(),
    ]:
        indices = pa.array([0, 1, 2, 0, 1, 2], data_type)
        dict_arr = pa.DictionaryArray.from_arrays(indices, dictionary)
        round_tripped = round_trip(dict_arr)
        assert round_tripped == dict_arr
        assert round_tripped.type == dict_arr.type


def test_empty_structs(tmp_path):
    schema = pa.schema([pa.field("empties", pa.struct([]))])
    table = pa.table({"empties": [{}] * 3}, schema=schema)
    path = tmp_path / "foo.lance"
    with LanceFileWriter(str(path)) as writer:
        writer.write_batch(table)
    reader = LanceFileReader(str(path))
    round_tripped = reader.read_all().to_table()
    assert round_tripped == table


def test_write_read_global_buffer(tmp_path):
    table = pa.table({"a": [1, 2, 3]})
    path = tmp_path / "foo.lance"
    global_buffer_text = "hello"
    global_buffer_bytes = bytes(global_buffer_text, "utf-8")
    with LanceFileWriter(str(path)) as writer:
        writer.write_batch(table)
        global_buffer_pos = writer.add_global_buffer(global_buffer_bytes)
    reader = LanceFileReader(str(path))
    assert reader.read_all().to_table() == table
    assert reader.metadata().global_buffers[global_buffer_pos].size == len(
        global_buffer_bytes
    )
    assert (
        bytes(reader.read_global_buffer(global_buffer_pos)).decode()
        == global_buffer_text
    )


def test_write_read_additional_schema_metadata(tmp_path):
    table = pa.table({"a": [1, 2, 3]})
    path = tmp_path / "foo.lance"
    schema_metadata_key = "foo"
    schema_metadata_value = "bar"
    with LanceFileWriter(str(path)) as writer:
        writer.write_batch(table)
        writer.add_schema_metadata(schema_metadata_key, schema_metadata_value)
    reader = LanceFileReader(str(path))
    assert reader.read_all().to_table() == table
    assert (
        reader.metadata().schema.metadata.get(schema_metadata_key.encode()).decode()
        == schema_metadata_value
    )


def test_writer_maintains_order(tmp_path):
    # 100Ki strings, each string is a couple of KiBs
    big_strings = [f"{i}" * 1024 for i in range(100 * 1024)]
    table = pa.table({"big_strings": big_strings})

    for i in range(4):
        path = tmp_path / f"foo-{i}.lance"
        with LanceFileWriter(str(path)) as writer:
            writer.write_batch(table)

        reader = LanceFileReader(str(path))
        result = reader.read_all().to_table()
        assert result == table


def test_compression(tmp_path):
    # 10Ki strings, which should be highly compressible, but not eligible for dictionary
    compressible_strings = [f"compress_me_please-{i}" for i in range(10 * 1024)]
    table_default = pa.table({"compressible_strings": compressible_strings})

    schema_compress = pa.schema(
        [
            pa.field(
                "compressible_strings",
                pa.string(),
                metadata={"lance-encoding:compression": "zstd"},
            )
        ]
    )
    table_compress = pa.table(
        {"compressible_strings": compressible_strings}, schema=schema_compress
    )

    with LanceFileWriter(str(tmp_path / "default.lance")) as writer:
        writer.write_batch(table_default)

    with LanceFileWriter(str(tmp_path / "compress.lance"), schema_compress) as writer:
        writer.write_batch(table_compress)

    size_default = os.path.getsize(tmp_path / "default.lance")
    size_compress = os.path.getsize(tmp_path / "compress.lance")

    assert size_compress < size_default


def test_blob(tmp_path):
    # 100 1MiB values.  If we store as regular large_binary we end up
    # with several pages of values.  If we store as a blob we get a
    # single page
    expected = pa.table(
        {"val": pa.array([b"0" * (1024 * 1024)] * 100, pa.large_binary())}
    )
    schema_no_blob = pa.schema([pa.field("val", pa.large_binary())])
    schema_blob = pa.schema(
        [pa.field("val", pa.large_binary(), metadata={"lance-encoding:blob": "true"})]
    )

    path = tmp_path / "no_blob.lance"
    with LanceFileWriter(str(path), schema_no_blob) as writer:
        for _ in range(100):
            vals = pa.array([b"0" * (1024 * 1024)], pa.large_binary())
            writer.write_batch(pa.table({"val": vals}))

    reader = LanceFileReader(str(path))
    assert len(reader.metadata().columns[0].pages) > 1
    assert reader.read_all().to_table() == expected

    path = tmp_path / "blob.lance"
    with LanceFileWriter(str(path), schema_blob) as writer:
        for _ in range(100):
            vals = pa.array([b"0" * (1024 * 1024)], pa.large_binary())
            writer.write_batch(pa.table({"val": vals}))

    reader = LanceFileReader(str(path))
    assert len(reader.metadata().columns[0].pages) == 1

    actual = reader.read_all().to_table()

    assert actual.num_rows == expected.num_rows
    for row_num in range(expected.num_rows):
        actual_bytes = actual.column("val").chunk(0)[row_num].as_py()
        expected_bytes = expected.column("val").chunk(0)[row_num].as_py()
        assert len(actual_bytes) == len(expected_bytes)
        assert actual_bytes == expected_bytes


def test_multithreaded_writer(tmp_path):
    """Test concurrent multi-threaded writing to the same LanceFileWriter"""
    path = tmp_path / "multithreaded.lance"
    schema = pa.schema(
        [
            pa.field("thread_id", pa.int64()),
            pa.field("value", pa.int64()),
            pa.field("data", pa.string()),
        ]
    )

    # Used to store all written data for subsequent validation
    all_data = []
    data_lock = threading.Lock()

    def write_thread_data(thread_id, writer, num_records):
        """Function for individual thread to write data"""
        thread_data = []
        for i in range(num_records):
            record = {
                "thread_id": thread_id,
                "value": thread_id * 1000 + i,
                "data": f"thread_{thread_id}_record_{i}",
            }
            thread_data.append(record)

        # Create pyarrow table
        table = pa.table(
            {
                "thread_id": [r["thread_id"] for r in thread_data],
                "value": [r["value"] for r in thread_data],
                "data": [r["data"] for r in thread_data],
            }
        )

        # Write data
        writer.write_batch(table)

        # Record written data for validation
        with data_lock:
            all_data.extend(thread_data)

    # Test parameters
    num_threads = 5
    records_per_thread = 100

    # Create writer and start multi-threaded writing
    with LanceFileWriter(str(path), schema) as writer:
        threads = []

        # Start multiple threads
        for thread_id in range(num_threads):
            thread = threading.Thread(
                target=write_thread_data, args=(thread_id, writer, records_per_thread)
            )
            threads.append(thread)
            thread.start()

        # Wait for all threads to complete
        for thread in threads:
            thread.join()

    # Validate written data
    reader = LanceFileReader(str(path))
    result_table = reader.read_all().to_table()

    # Check if total row count is correct
    expected_total_rows = num_threads * records_per_thread
    assert result_table.num_rows == expected_total_rows, (
        f"Expected {expected_total_rows} rows, got {result_table.num_rows}"
    )

    # Check data content correctness (order may differ, but data should be complete)
    # Convert results to dictionary list for comparison
    result_data = [
        {
            "thread_id": result_table.column("thread_id").chunk(0)[i].as_py(),
            "value": result_table.column("value").chunk(0)[i].as_py(),
            "data": result_table.column("data").chunk(0)[i].as_py(),
        }
        for i in range(result_table.num_rows)
    ]

    # Verify all data exists (order not considered)
    all_data_sorted = sorted(all_data, key=lambda x: (x["thread_id"], x["value"]))
    result_data_sorted = sorted(result_data, key=lambda x: (x["thread_id"], x["value"]))

    assert len(all_data_sorted) == len(result_data_sorted)

    # Compare data item by item
    for expected, actual in zip(all_data_sorted, result_data_sorted):
        assert expected == actual, f"Data mismatch: expected {expected}, got {actual}"

    # Verify data from each thread exists
    thread_ids_in_result = set(result_table.column("thread_id").chunk(0).to_pylist())
    expected_thread_ids = set(range(num_threads))
    assert thread_ids_in_result == expected_thread_ids

    # Verify record count for each thread
    for thread_id in range(num_threads):
        thread_rows = result_table.filter(
            pc.equal(result_table.column("thread_id"), thread_id)
        )
        assert thread_rows.num_rows == records_per_thread


def test_session_list_all_files(tmp_path):
    """Test that LanceFileSession.list() returns all files with relative paths"""
    session = LanceFileSession(str(tmp_path))
    schema = pa.schema([pa.field("x", pa.int64())])

    # Write files at different levels
    with session.open_writer("file1.lance", schema=schema) as writer:
        writer.write_batch(pa.table({"x": [1]}))

    with session.open_writer("file2.lance", schema=schema) as writer:
        writer.write_batch(pa.table({"x": [2]}))

    with session.open_writer("subdir/file3.lance", schema=schema) as writer:
        writer.write_batch(pa.table({"x": [3]}))

    with session.open_writer("subdir/file4.lance", schema=schema) as writer:
        writer.write_batch(pa.table({"x": [4]}))

    with session.open_writer("other/file5.lance", schema=schema) as writer:
        writer.write_batch(pa.table({"x": [5]}))

    # List all files
    files = sorted(session.list())

    # Verify relative paths (no absolute paths)
    assert files == [
        "file1.lance",
        "file2.lance",
        "other/file5.lance",
        "subdir/file3.lance",
        "subdir/file4.lance",
    ]

    # Verify no absolute paths
    for f in files:
        assert not f.startswith("/")
        assert str(tmp_path) not in f


def test_session_list_with_prefix(tmp_path):
    """Test that LanceFileSession.list() filters by prefix correctly"""
    session = LanceFileSession(str(tmp_path))
    schema = pa.schema([pa.field("x", pa.int64())])

    # Write files in different directories
    with session.open_writer("file1.lance", schema=schema) as writer:
        writer.write_batch(pa.table({"x": [1]}))

    with session.open_writer("subdir/file2.lance", schema=schema) as writer:
        writer.write_batch(pa.table({"x": [2]}))

    with session.open_writer("subdir/file3.lance", schema=schema) as writer:
        writer.write_batch(pa.table({"x": [3]}))

    with session.open_writer("other/file4.lance", schema=schema) as writer:
        writer.write_batch(pa.table({"x": [4]}))

    # List with prefix "subdir"
    subdir_files = sorted(session.list("subdir"))
    assert subdir_files == ["subdir/file2.lance", "subdir/file3.lance"]

    # List with prefix "other"
    other_files = sorted(session.list("other"))
    assert other_files == ["other/file4.lance"]

    # List with non-existent prefix
    empty = session.list("nonexistent")
    assert empty == []


def test_session_list_with_trailing_slash(tmp_path):
    """Test that LanceFileSession.list() handles trailing slashes correctly"""
    session = LanceFileSession(str(tmp_path))
    schema = pa.schema([pa.field("x", pa.int64())])

    with session.open_writer("dir/file.lance", schema=schema) as writer:
        writer.write_batch(pa.table({"x": [1]}))

    # Both with and without trailing slash should work
    files_no_slash = session.list("dir")
    files_with_slash = session.list("dir/")

    assert files_no_slash == files_with_slash
    assert files_no_slash == ["dir/file.lance"]


def test_session_contains(tmp_path):
    """Test that LanceFileSession.contains() works correctly"""
    session = LanceFileSession(str(tmp_path))
    schema = pa.schema([pa.field("x", pa.int64())])

    # File doesn't exist yet
    assert not session.contains("test.lance")

    # Write a file
    with session.open_writer("test.lance", schema=schema) as writer:
        writer.write_batch(pa.table({"x": [1]}))

    # File exists now
    assert session.contains("test.lance")

    # Nested file
    with session.open_writer("subdir/nested.lance", schema=schema) as writer:
        writer.write_batch(pa.table({"x": [2]}))

    assert session.contains("subdir/nested.lance")
    assert not session.contains("subdir/nonexistent.lance")
